import pandas as pd
import matplotlib.pyplot as plt

# ==================== 配置参数 ====================
INPUT_FILE = "ori_message.xlsx"  # 原始数据文件路径
OUTPUT_FILE = "monthly_average_data.xlsx"  # 输出文件路径


# ==================== 数据读取与预处理 ====================
def process_data():
    excel_file = pd.ExcelFile(INPUT_FILE)
    years = ['2016', '2017', '2018', '2019', '2020', '2021']  # 对应工作表名称
    all_monthly_data = []
    for year in years:
        df = excel_file.parse(year)
        df['年'] = df['年'].ffill().astype(int)
        df['月'] = df['月'].ffill().astype(int)
        grouped = df.groupby(['年', '月']).agg({
            '流量(m3/s)': 'mean',
            '含沙量(kg/m3) ': 'mean'
        }).reset_index()
        grouped['排沙量(kg/s)'] = grouped['流量(m3/s)'] * grouped['含沙量(kg/m3) ']
        grouped = grouped.round({
            '流量(m3/s)': 2,
            '含沙量(kg/m3) ': 4,
            '排沙量(kg/s)': 2
        })
        all_monthly_data.append(grouped)
    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)
    result = result.rename(columns={'含沙量(kg/m3) ': '含沙量(kg/m3)'})
    result.to_excel(OUTPUT_FILE, index=False)
    return result


# ==================== 绘制5-10月流量走势图 ====================
def plot_5_10_flow(data):
    # 筛选5-10月数据
    filtered_data = data[data['月'].between(5, 10)]

    # 定义年份顺序和颜色（与图片图例顺序一致：2016到2021年）
    years_order = [2016, 2017, 2018, 2019, 2020, 2021]
    colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']  # 不同颜色区分年份

    plt.figure(figsize=(10, 6))

    for i, year in enumerate(years_order):
        year_data = filtered_data[filtered_data['年'] == year]
        # 提取月份（5-10月，转换为中文标签）
        months = year_data['月'].map({5: '5月', 6: '6月', 7: '7月', 8: '8月', 9: '9月', 10: '10月'})
        plt.plot(months, year_data['流量(m3/s)'], marker='o', linestyle='-',
                 label=str(year), color=colors[i], linewidth=1.5)

    # 图表设置（严格匹配图片描述）
    plt.title('六年中5-10月的月均流量走势图', fontsize=12)
    plt.xlabel('')  # 图片中x轴无标签，仅显示月份
    plt.ylabel('月平均流量 (m³/s)', fontsize=10)
    plt.xticks(rotation=0)  # 月份标签不旋转
    plt.legend(loc='best', title='年份', bbox_to_anchor=(1, 1))  # 图例位于右上角外部
    plt.grid(linestyle='--', alpha=0.5, axis='y')  # 仅显示y轴网格

    # 调整布局，避免图例遮挡
    plt.tight_layout(rect=[0, 0, 0.85, 1])

    # 保存图片（严格匹配文件名）
    plt.savefig('图29.六年中5-10月的月均流量走势图.png', dpi=300, bbox_inches='tight')
    plt.show()


# ==================== 主程序 ====================
if __name__ == "__main__":
    processed_data = process_data()
    plot_5_10_flow(processed_data)